// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/exceptions.h>
#include <pollux/exec/window_function.h>
#include <pollux/expression/function_signature.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::functions::window {
    // Types of rank functions.
    enum class RankType {
        kRank,
        kDenseRank,
        kPercentRank,
    };

    namespace {
        template<RankType TRank, typename TResult>
        class RankFunction : public exec::WindowFunction {
        public:
            explicit RankFunction(const TypePtr &resultType)
                : WindowFunction(resultType, nullptr, nullptr) {
            }

            void resetPartition(const exec::WindowPartition *partition) override {
                rank_ = 1;
                currentPeerGroupStart_ = 0;
                previousPeerCount_ = 0;
                numPartitionRows_ = partition->numRows();
            }

            void apply(
                const BufferPtr &peerGroupStarts,
                const BufferPtr & /*peerGroupEnds*/,
                const BufferPtr & /*frameStarts*/,
                const BufferPtr & /*frameEnds*/,
                const SelectivityVector &validRows,
                vector_size_t resultOffset,
                const VectorPtr &result) override {
                int numRows = peerGroupStarts->size() / sizeof(vector_size_t);
                auto *rawPeerStarts = peerGroupStarts->as<vector_size_t>();
                auto rawValues = result->as_flat_vector<TResult>()->mutableRawValues();

                for (int i = 0; i < numRows; i++) {
                    auto start = rawPeerStarts[i];
                    if (start != currentPeerGroupStart_) {
                        currentPeerGroupStart_ = start;
                        if constexpr (TRank == RankType::kDenseRank) {
                            if (previousPeerCount_ > 0) {
                                ++rank_;
                            }
                        } else {
                            rank_ += previousPeerCount_;
                        }
                        previousPeerCount_ = 0;
                    }

                    if constexpr (TRank == RankType::kPercentRank) {
                        rawValues[resultOffset + i] = (numPartitionRows_ == 1)
                                                          ? 0
                                                          : double(rank_ - 1) / (numPartitionRows_ - 1);
                    } else {
                        rawValues[resultOffset + i] = rank_;
                    }
                    previousPeerCount_ += 1;
                }
            }

        private:
            int32_t currentPeerGroupStart_ = 0;
            int32_t previousPeerCount_ = 0;
            int64_t rank_ = 1;
            vector_size_t numPartitionRows_ = 1;
        };
    } // namespace

    template<RankType TRank, typename TResult>
    void registerRankInternal(
        const std::string &name,
        const std::string &returnType) {
        std::vector<exec::FunctionSignaturePtr> signatures{
            exec::FunctionSignatureBuilder().returnType(returnType).build(),
        };

        auto windowFunctionFactory =
                [name](
            const std::vector<exec::WindowFunctionArg> & /*args*/,
            const TypePtr &resultType,
            bool /*ignoreNulls*/,
            pollux::memory::MemoryPool * /*pool*/,
            HashStringAllocator * /*stringAllocator*/,
            const core::QueryConfig & /*queryConfig*/)
            -> std::unique_ptr<exec::WindowFunction> {
            return std::make_unique<RankFunction<TRank, TResult> >(resultType);
        };

        if constexpr (TRank == RankType::kRank || TRank == RankType::kDenseRank) {
            exec::registerWindowFunction(
                name,
                std::move(signatures),
                {exec::WindowFunction::ProcessMode::kRows, false},
                std::move(windowFunctionFactory));
        } else {
            exec::registerWindowFunction(
                name,
                std::move(signatures),
                exec::WindowFunction::Metadata::defaultMetadata(),
                std::move(windowFunctionFactory));
        }
    }

    void registerRankBigint(const std::string &name) {
        registerRankInternal<RankType::kRank, int64_t>(name, "bigint");
    }

    void registerRankInteger(const std::string &name) {
        registerRankInternal<RankType::kRank, int32_t>(name, "integer");
    }

    void registerDenseRankBigint(const std::string &name) {
        registerRankInternal<RankType::kDenseRank, int64_t>(name, "bigint");
    }

    void registerDenseRankInteger(const std::string &name) {
        registerRankInternal<RankType::kDenseRank, int32_t>(name, "integer");
    }

    void registerPercentRank(const std::string &name) {
        registerRankInternal<RankType::kPercentRank, double>(name, "double");
    }
} // namespace kumo::pollux::functions::window
